import socket
import ssl
import http.client
import re
import urllib.parse
from .logger import get_logger

logger = get_logger()

def detect_service(target, port, timeout=2):
    """Detect service details for an open port"""
    service_info = {
        'name': 'unknown',
        'version': None,
        'product': None,
        'http_code': None,
        'http_title': None,
        'http_server': None,
        'ssl': False,
        'banner': None
    }
    
    try:
        # Try SSL connection first
        try:
            context = ssl.create_default_context()
            context.check_hostname = False
            context.verify_mode = ssl.CERT_NONE
            with socket.create_connection((target, port), timeout=timeout) as sock:
                with context.wrap_socket(sock, server_hostname=target) as ssock:
                    service_info['ssl'] = True
                    cert = ssock.getpeercert(binary_form=True)
                    if cert:
                        service_info['version'] = 'TLS'
        except:
            pass

        # Check if it's HTTP(S)
        if port in [80, 443, 8080, 8443] or service_info['ssl']:
            protocol = 'https' if (port == 443 or service_info['ssl']) else 'http'
            service_info.update(get_http_info(target, port, protocol))
            return service_info
            
        # Check if it's FTP
        if port == 21:
            service_info.update(get_ftp_info(target, port, timeout))
            return service_info

        # Try banner grabbing for other services
        banner = grab_banner(target, port, timeout)
        if banner:
            service_info['banner'] = banner
            service_info.update(parse_banner(banner))

        # Get service name from system database
        try:
            service_info['name'] = socket.getservbyport(port)
        except:
            pass

    except Exception as e:
        logger.debug(f"Error detecting service on {target}:{port} - {str(e)}")

    return service_info

def get_http_info(target, port, protocol):
    """Get HTTP service information"""
    info = {
        'name': 'http',
        'http_code': None,
        'http_title': None,
        'http_server': None
    }
    
    try:
        conn = http.client.HTTPSConnection(target, port, timeout=2) if protocol == 'https' \
              else http.client.HTTPConnection(target, port, timeout=2)
        
        conn.request('GET', '/')
        response = conn.getresponse()
        
        info['http_code'] = response.status
        info['http_server'] = response.getheader('Server')
        
        # Get page title
        body = response.read().decode('utf-8', errors='ignore')
        title_match = re.search('<title>(.*?)</title>', body, re.IGNORECASE)
        if title_match:
            info['http_title'] = title_match.group(1).strip()
        
        conn.close()
    except Exception as e:
        logger.debug(f"Error getting HTTP info: {str(e)}")
    
    return info

def get_ftp_info(target, port, timeout):
    """Get FTP service information"""
    info = {
        'name': 'ftp',
        'banner': None,
        'version': None,
        'product': None,
        'anonymous': False
    }
    
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.settimeout(timeout)
            sock.connect((target, port))
            
            # 读取欢迎banner
            banner = sock.recv(1024).decode('utf-8', errors='ignore').strip()
            info['banner'] = banner
            info.update(parse_banner(banner))
            
            # 尝试匿名登录
            try:
                sock.send(b'USER anonymous\r\n')
                resp = sock.recv(1024).decode('utf-8', errors='ignore')
                if '331' in resp:  # 331 = User name okay, need password
                    sock.send(b'PASS anonymous@example.com\r\n')
                    resp = sock.recv(1024).decode('utf-8', errors='ignore')
                    info['anonymous'] = '230' in resp  # 230 = User logged in
            except:
                pass
                
            # 获取系统信息
            try:
                sock.send(b'SYST\r\n')
                resp = sock.recv(1024).decode('utf-8', errors='ignore')
                if '215' in resp:  # 215 = System information
                    info['system'] = resp.split('215 ')[1].strip()
            except:
                pass
                
            # 获取特性列表
            try:
                sock.send(b'FEAT\r\n')
                resp = sock.recv(1024).decode('utf-8', errors='ignore')
                if '211-' in resp:  # 211 = System status
                    features = []
                    for line in resp.split('\n'):
                        if line.startswith(' '):
                            features.append(line.strip())
                    if features:
                        info['features'] = features
            except:
                pass
            
    except Exception as e:
        logger.debug(f"Error getting FTP info: {str(e)}")
    
    return info

def grab_banner(target, port, timeout):
    """Grab service banner"""
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.settimeout(timeout)
            sock.connect((target, port))
            # 对于大多数服务，不需要发送数据就能收到banner
            banner = sock.recv(1024).decode('utf-8', errors='ignore')
            if not banner:
                # 如果没有收到banner，尝试发送一个换行符
                sock.send(b'\r\n')
                banner = sock.recv(1024).decode('utf-8', errors='ignore')
            return banner.strip()
    except:
        return None

def parse_banner(banner):
    """Parse service banner for version information"""
    info = {
        'version': None,
        'product': None
    }
    
    if not banner:
        return info

    # Common version patterns
    version_patterns = [
        r'(?i)version[\s:]+([\\d.]+)',
        r'(?i)([\\w-]+)(?:[\s/]+)(\\d+(?:\\.\\d+)+)',
        r'(?i)(\\w+)[\\s/]+(\\d+(?:\\.\\d+)+)'
    ]

    for pattern in version_patterns:
        match = re.search(pattern, banner)
        if match:
            info['product'] = match.group(1)
            info['version'] = match.group(2)
            break

    return info 